# -*- coding: utf-8 -*-
import json
import random
from openai import OpenAI

class HSKService:
    def __init__(self):
        # 豆包
        # self.client = OpenAI(
        #     api_key="a2f51613-131e-4ded-a077-c087b2ea3f4e",
        #     base_url="https://ark.cn-beijing.volces.com/api/v3",
        # )
        self.client = OpenAI(
            api_key="sk-a9ff8d1c5e8946d6a30dba8bc16a1220",
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )
        self.data = self._load_word_data("./data/hsk_topic_tag.json")
        self.hsk_level_data = self._load_word_data("./data/hsk_level.json")
        
    def _load_word_data(self, filepath):
        """加载词汇数据"""
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                return json.load(f)
        except FileNotFoundError:
            raise Exception(f"数据文件未找到: {filepath}")
        except json.JSONDecodeError:
            raise Exception(f"数据文件格式错误: {filepath}")

    def _extract_unique_labels(self):
        """提取唯一标签"""
        tags = set()
        for tag_list in self.data.values():
            for tag in tag_list:
                tags.add(tag)
        return list(tags)

    def _match_labels_with_llm(self, user_input, labels):
        """使用LLM匹配标签"""
        prompt = f"""
            你是一个中文语义匹配助手。
            用户输入的主题是：【{user_input}】
            下面是可选的标签类别：
            {labels}
            请根据意思相近度，帮我从列表里选择 1~5 个最相关的标签。
            只返回用中文逗号分隔的标签字符串，不要解释。
        """
        try:
            response = self.client.chat.completions.create(
                # model="doubao-seed-1-6-250615",
                model = "qwen-plus",
                messages=[
                    {"role": "system", "content": "你是一个高效的语义匹配助手，只返回标签，不解释。"},
                    {"role": "user", "content": prompt}
                ]
            )
            matched = response.choices[0].message.content
            matched_labels = [label.strip() for label in matched.split(",") if label.strip()]
            return matched_labels
        except Exception as e:
            raise Exception(f"LLM标签匹配失败: {str(e)}")

    def _filter_words_by_labels(self, matched_labels):
        """根据标签过滤词汇"""
        matched_words = [
            word for word, tags in self.data.items()
            if any(label in tags for label in matched_labels)
        ]
        return matched_words

    def _filter_words_by_hsk_level(self, words, level):
        """根据HSK等级过滤词汇"""
        result = []
        low_result = []
        high_result = []
        
        for word in words:
            if word not in self.hsk_level_data:
                continue
            word_level = self.hsk_level_data[word]
            if word_level == level:
                result.append(word)
            elif word_level < level:
                low_result.append(word)
            elif word_level > level:
                high_result.append(word)
        return result, low_result, high_result

    def _judge_word_relevance(self, theme, words, level):
        map_hsk = {
            1: '小学低年级', 2: '小学中年级', 3: '小学高年级',
            4: '初中学生', 5: '高中学生', 6: '大学生',
            7: '研究生', 8: '博士生', 9: '博士后',
        }

        prompt = f"""
        你是一个中文语义匹配助手。

        用户输入的主题是：【{theme}】

        以下是用户提供的一组词汇：
        {words}

        你的任务包括两个部分：
        1. 请你根据每个词汇与主题的语义相关性，对每个词进行打分（范围为 0 到 100）。
        2. 请你补充 5 个与该主题高度相关的词汇，这些词汇应适合 {map_hsk[level]} 水平的中文学习者，并为每个补充词汇也打分（相关性评分范围同样是 0 到 100）。

        ⚠️ 请你严格只返回以下格式的 JSON 字符串，不添加任何注释、说明、代码块标记（如 ```json）或多余文字：

        {{
          "scores": {{
            "原始词汇1": 分数,
            "原始词汇2": 分数
            // ...
          }},
          "supplement_scores": {{
            "补充词汇1": 分数,
            "补充词汇2": 分数
            // ...
          }}
        }}

        返回结果必须以左大括号 {{ 开头，以右大括号 }} 结尾，符合 JSON 语法规范，确保可以直接使用 Python 的 json.loads() 成功解析。
        """

        try:
            response = self.client.chat.completions.create(
                # model="doubao-seed-1-6-250615",
                model="qwen-plus",
                messages=[
                    {"role": "system", "content": "你是一个中文语义匹配助手。只返回标签，不解释。"},
                    {"role": "user", "content": prompt}
                ]
            )

            content = response.choices[0].message.content.strip()

            # 解析为 JSON
            result = json.loads(content)

            # 简单验证结构是否符合预期
            if not isinstance(result, dict) or "scores" not in result or "supplement_scores" not in result:
                raise ValueError("返回结果不符合预期格式")

            return result  # 返回包含 scores 和 supplement_words 的 dict

        except json.JSONDecodeError as e:
            raise Exception(f"解析 JSON 失败: {e}\n返回内容: {content}")
        except Exception as e:
            raise Exception(f"LLM标签匹配失败: {str(e)}")
    def get_theme_labels(self, user_input):
        """Step1: 获取主题标签"""
        if not user_input or not user_input.strip():
            raise ValueError("用户输入不能为空")
        
        labels = self._extract_unique_labels()
        matched_labels = self._match_labels_with_llm(user_input, labels)
        return matched_labels

    def get_words_sample(self, user_input, level=3, sample_size=10):
        """Step2: 获取词汇样本"""
        if not user_input or not user_input.strip():
            raise ValueError("用户输入不能为空")
        if not isinstance(level, int) or level < 1 or level > 9:
            raise ValueError("HSK等级必须是1-9之间的整数")
        if not isinstance(sample_size, int) or sample_size < 1:
            raise ValueError("样本大小必须是正整数")
        matched_labels = self.get_theme_labels(user_input)
        matched_words = self._filter_words_by_labels(matched_labels)
        if not matched_words:
            # 如果没有匹配的词汇，随机选择
            matched_words = list(self.data.keys())
        matched_words, low_result, high_result = self._filter_words_by_hsk_level(matched_words, level)
        if not matched_words:
            raise Exception(f"未找到HSK{level}级别的相关词汇")
        if len(matched_words) < sample_size:
            needed = sample_size - len(matched_words)
            # 优先从 low_result 中补充
            supplement = low_result[:needed]
            # 如果还不够，再从 high_result 中补充
            if len(supplement) < needed:
                supplement += high_result[:(needed - len(supplement))]
            matched_words += supplement
        # 评价语义相关性，并增加补充词汇及其相关性
        result = self._judge_word_relevance(user_input, matched_words, level)
        return result




    def generate_article(self, topic, hsk_level=3, article_len=300, words=None, other_requirements=""):
        """Step3: 生成文章"""
        if not topic or not topic.strip():
            raise ValueError("主题不能为空")
            
        if words is None:
            words = []
            
        map_hsk = {
            1: '小学低年级', 2: '小学中年级', 3: '小学高年级',
            4: '初中学生', 5: '高中学生', 6: '大学生',
            7: '研究生', 8: '博士生', 9: '博士后',
        }
        
        if hsk_level not in map_hsk:
            raise ValueError("HSK等级必须是1-9之间的整数")
        
        words_str = ', '.join(words) if words else ''
        prompt = f'''请你根据以下要求写3篇中文短文：
            主题：{topic}
            读者：写给正在读{map_hsk[hsk_level]}
            难度：HSK {hsk_level}级
            目标长度：约{article_len}字
            必须包含词语：{words_str}
            其他要求：{other_requirements}
            
            写作要求：
            要求内容自然、表达清晰，并在文中自然地使用所有指定词语。
            返回格式：```json{{"output": ["...", "...", "..."]}}```'''

        try:
            completion = self.client.chat.completions.create(
                # model='doubao-seed-1-6-250615',
                model="qwen-plus",
                messages=[
                    {
                        "role": "system",
                        "content": "你是一个国际中文教育助手，擅长根据指定词语、难度等级和长度要求，生成符合教学需求的中文短文。"
                    },
                    {"role": "user", "content": prompt}
                ],
                extra_body={
                    "enable_enhancement": True,
                },

            )
            
            ret = completion.choices[0].message.content
            
            # 处理返回的JSON格式
            if ret.startswith('```json'):
                ret = ret[7:]
            if ret.endswith('```'):
                ret = ret[:-3]
            
            ret_data = json.loads(ret)
            return ret_data.get('output', [])
            
        except json.JSONDecodeError:
            raise Exception("LLM返回的内容不是有效的JSON格式")
        except Exception as e:
            raise Exception(f"文章生成失败: {str(e)}")

    def generate_article_stream(self, topic, sid,hsk_level=3, article_len=300, words=None, other_requirements=""):
        print(f"调用 generate_article_stream 时的 sid: {sid}")
        """
        生成文章（流式），每次产出一小段内容通过 socket 推送
        """
        if not topic or not topic.strip():
            raise ValueError("主题不能为空")

        if words is None:
            words = []

        map_hsk = {
            1: '小学低年级', 2: '小学中年级', 3: '小学高年级',
            4: '初中学生', 5: '高中学生', 6: '大学生',
            7: '研究生', 8: '博士生', 9: '博士后',
        }

        if hsk_level not in map_hsk:
            raise ValueError("HSK等级必须是1-9之间的整数")

        words_str = ', '.join(words) if words else ''
        prompt = f'''请你根据以下要求写1篇中文短文：
        主题：{topic}
        读者：写给正在读{map_hsk[hsk_level]}
        难度：HSK {hsk_level}级
        目标长度：约{article_len}字
        必须包含词语：{words_str}
        其他要求：{other_requirements}
    
        写作要求：
        要求内容自然、表达清晰，并在文中自然地使用所有指定词语。
        直接开始写作正文，不需要返回 JSON 格式，不要包含任何格式标记。'''

        try:
            stream = self.client.chat.completions.create(
                model="qwen-plus",
                messages=[
                    {"role": "system", "content": "你是一个国际中文教育助手，擅长生成符合教学要求的中文短文。"},
                    {"role": "user", "content": prompt}
                ],
                stream=True
            )
            result = ""
            for chunk in stream:
                if chunk.choices and chunk.choices[0].delta:
                    content_piece = chunk.choices[0].delta.content if chunk.choices[0].delta.content else ""
                    result += content_piece
                    print(f"推送内容：{content_piece}")
                    # 👇 如果启用了 WebSocket 推送，逐段推送
                    print(sid)
                    if sid:
                        from task import send_to_socketio
                        send_to_socketio(sid, step=3, result=content_piece)  # 只发这段，不是整篇
            return result  # 返回整篇文章

        except Exception as e:
            raise Exception(f"生成短文失败: {str(e)}")


if __name__ == "__main__":
    # 使用示例
    service = HSKService()
    user_input = input("请输入你的主题词：")  # 例如：旅游
    user_level = input("请输入hsk等级：")
    hsk_words = service.get_words_sample(user_input, int(user_level), 5)
    combined_scores = {**hsk_words["scores"], **hsk_words["supplement_scores"]}
    # 按分数降序排序
    sorted_words = sorted(combined_scores.items(), key=lambda x: x[1], reverse=True)
    top_words = [word for word, score in sorted_words[:10]]  # 提取前10个单词（去掉分数）

    print("\n=== 抽样结果 ===")
    print(hsk_words)
    print("\n=== 最终抽样结果 ===")
    print(top_words)
    print(service.generate_article(
        topic=user_input,
        hsk_level=int(user_level),
        article_len=300,
        words=top_words  # 传入纯单词列表，而不是元组列表
    ))
